# -*- coding: utf-8 -*-
"""
Created on Sun Oct 11 10:09:35 2020
用于库伦耦合的半导体双量子点模型
128初始态对制备保真度

@author: Waikikilick
"""

import numpy as np
from scipy.linalg import expm
from time import *

import copy
import multiprocessing as mp
np.random.seed(1)
I = np.matrix(np.identity(2, dtype=complex))
sx = np.mat([[0, 1], [1, 0]], dtype=complex)
sy = np.mat([[0, -1j], [1j, 0]], dtype=complex)
sz = np.mat([[1, 0], [0, -1]], dtype=complex)
h_1 = 1
h_2 = 1
coupling = 1/2
T = 10*np.pi
dt = np.pi/2
step_max = T/dt
# init_psi = np.mat([[1], [0], [0], [0]], dtype=complex)
# target_psi = np.mat([[1], [0], [0], [1]], dtype=complex)/np.sqrt(2) #最终目标态
init_set_num = 512
target_set_num = init_set_num

action_space =np.array([[1,1],
                        [1,2],
                        [1,3],
                        [1,4],
                        [1,5],
                        [2,1],
                        [2,2],
                        [2,3],
                        [2,4],
                        [2,5],
                        [3,1],
                        [3,2],
                        [3,3],
                        [3,4],
                        [3,5],
                        [4,1],
                        [4,2],
                        [4,3],
                        [4,4],
                        [4,5],
                        [5,1],
                        [5,2],
                        [5,3],
                        [5,4],
                        [5,5]
                        ])
n_actions = len(action_space)

noise_normal = np.array([ 1.62434536, -0.61175641, -0.52817175, -1.07296862,  0.86540763,
                         -2.3015387 ,  1.74481176, -0.7612069 ,  0.3190391 , -0.24937038,
                          1.46210794, -2.06014071, -0.3224172 , -0.38405435,  1.13376944,
                         -1.09989127, -0.17242821, -0.87785842,  0.04221375,  0.58281521,
                         -1.10061918,  1.14472371,  0.90159072,  0.50249434,  0.90085595,
                         -0.68372786, -0.12289023, -0.93576943, -0.26788808,  0.53035547,
                         -0.69166075, -0.39675353, -0.6871727 , -0.84520564, -0.67124613,
                         -0.0126646 , -1.11731035,  0.2344157 ,  1.65980218,  0.74204416,
                         -0.19183555,-0.88762896,  -0.74715829,  1.6924546 ,  0.05080775,
                         -0.63699565,  0.19091548,  2.10025514,  0.12015895,  0.61720311,
                         0.30017032, -0.35224985, -1.1425182,  -0.34934272,  -0.20889423,
                         0.58662319,  0.83898341,  0.93110208,  0.28558733,  0.88514116, 
                         -0.75439794,  1.25286816,  0.51292982, -0.29809284,  0.48851815,
                         -0.07557171,  1.13162939,  1.51981682,  2.18557541, -1.39649634,
                         -1.44411381, -0.50446586,  0.16003707,  0.87616892,  0.31563495,
                         -2.02220122, -0.30620401,  0.82797464,  0.23009474,  0.76201118, 
                         -0.22232814, -0.20075807,  0.18656139,  0.41005165,  0.19829972,
                         0.11900865, -0.67066229,  0.37756379,  0.12182127,  1.12948391,
                         1.19891788,  0.18515642, -0.37528495, -0.63873041,  0.42349435,
                         0.07734007, -0.34385368,  0.04359686, -0.62000084,  0.69803203])
#noise_normal 为均值为 0 ，标准差为 1 的正态分布的随机数组成的数组
#该随机数由 np.random.seed(1) 生成: np.random.seed(1) \ noise_uniform = np.random.normal(loc=0.0, scale=1.0, size=41)
noise_drift = np.array([1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.])

alpha_num = 4
theta = [np.pi/8,np.pi/4,3*np.pi/8]
theta_1 = theta
theta_2 = theta
theta_3 = theta

alpha = np.linspace(0,np.pi*2,alpha_num,endpoint=False)
alpha_1 = alpha
alpha_2 = alpha
alpha_3 = alpha
alpha_4 = alpha

a_list_complex = np.matrix([[0,0,0,0]],dtype=complex) #第一行用来占位，否则无法和其他行并在一起，在最后要注意去掉这一行
for ii in range(3): #theta_1
    for jj in range(3): #theta_2
        for kk in range(3): #theta_3
            for mm in range(alpha_num): #alpha_1
                for nn in range(alpha_num): #alpha_2
                    for oo in range(alpha_num): #alpha_3
                        for pp in range(alpha_num): #alpha_4
                            
                            a_1_mo = np.cos(theta_1[ii])
                            a_2_mo = np.sin(theta_1[ii])*np.cos(theta_2[jj])
                            a_3_mo = np.sin(theta_1[ii])*np.sin(theta_2[jj])*np.cos(theta_3[kk])
                            a_4_mo = np.sin(theta_1[ii])*np.sin(theta_2[jj])*np.sin(theta_3[kk])
                            
                            a_1_real = a_1_mo*np.cos(alpha_1[mm])
                            a_1_imag = a_1_mo*np.sin(alpha_1[mm])
                            a_2_real = a_2_mo*np.cos(alpha_2[nn])
                            a_2_imag = a_2_mo*np.sin(alpha_2[nn])
                            a_3_real = a_3_mo*np.cos(alpha_3[oo])
                            a_3_imag = a_3_mo*np.sin(alpha_3[oo])
                            a_4_real = a_4_mo*np.cos(alpha_4[pp])
                            a_4_imag = a_4_mo*np.sin(alpha_4[pp])
                            
                            a_1_complex = a_1_real + a_1_imag*1j
                            a_2_complex = a_2_real + a_2_imag*1j
                            a_3_complex = a_3_real + a_3_imag*1j
                            a_4_complex = a_4_real + a_4_imag*1j
                            
                            a_complex = np.matrix([[ a_1_complex, a_2_complex, a_3_complex, a_4_complex]])
                            a_list_complex = np.row_stack((a_list_complex,a_complex))
                            
psi_set = np.array(np.delete(a_list_complex,0,axis=0)) # 删除矩阵的第一行
np.random.shuffle(psi_set) #打乱顺序
init_set = psi_set[:init_set_num]

def target_set():
    target_set = psi_set[init_set_num : init_set_num + target_set_num]
    return target_set

target_set = target_set()

def step_noise_J(psi,target_psi,action,noise_unit):
    J_1_noise, J_2_noise=  action_space[action,0]+noise_unit[0], action_space[action,1]+noise_unit[1]  # control field strength
    J_12_noise = J_1_noise * J_2_noise /2
    

    H =  (J_1_noise*np.kron(sz, I) + J_2_noise*np.kron(I, sz) + \
                    J_12_noise/2*np.kron((sz - I),(sz - I)) + \
       h_1*np.kron(sx,I) + h_2*np.kron(I,sx))*coupling
    U = expm(-1j * H * dt) 
    psi_ = U * psi
    fid = (np.abs(psi_.H * target_psi) ** 2).item(0).real 
    
    return fid, psi_

def step_noise_h(psi,target_psi,action,noise_unit):
    J_1, J_2=  action_space[action,0], action_space[action,1]  # control field strength
    J_12 = J_1 * J_2 /2
    h_1_noise, h_2_noise = h_1 + noise_unit[0], h_2 + noise_unit[1]

    H =  (J_1*np.kron(sz, I) + J_2*np.kron(I, sz) + \
                    J_12/2*np.kron((sz - I),(sz - I)) + \
       h_1_noise*np.kron(sx,I) + h_2_noise*np.kron(I,sx))*coupling
    U = expm(-1j * H * dt) 
    psi_ = U * psi
    fid = (np.abs(psi_.H * target_psi) ** 2).item(0).real   
    
    return fid, psi_

#动作直接选最优的
def step0(psi,target_psi,F):
    fid_list = []
    psi_list = []
    action_list = list(range(len(action_space)))
    for action in action_list:
        
        J_1, J_2=  action_space[action,0], action_space[action,1]  # control field strength
        J_12 = J_1 * J_2 /2
    
        H =  (J_1*np.kron(sz, I) + J_2*np.kron(I, sz) + \
                        J_12/2*np.kron((sz - I),(sz - I)) + \
           h_1*np.kron(sx,I) + h_2*np.kron(I,sx))*coupling
        U = expm(-1j * H * dt) 
        psi_ = U * psi
        fid = (np.abs(psi_.H * target_psi) ** 2).item(0).real 
        
        psi_list.append(psi_)
        fid_list.append(fid)
        best_action = fid_list.index(max(fid_list))
        best_fid = max(fid_list)
    psi_ = psi_list[best_action]
    # print(best_action)
    return best_action, best_fid, psi_


#动作选最优的，或者最差的
def step1(psi,target_psi,F):
    fid_list = []
    psi_list = []
    action_list = list(range(len(action_space)))
    for action in action_list:
        
        J_1, J_2=  action_space[action,0], action_space[action,1]  # control field strength
        J_12 = J_1 * J_2 /2
    
        H =  (J_1*np.kron(sz, I) + J_2*np.kron(I, sz) + \
                        J_12/2*np.kron((sz - I),(sz - I)) + \
           h_1*np.kron(sx,I) + h_2*np.kron(I,sx))*coupling

        U = expm(-1j * H * dt) 
        psi_ = U * psi
        fid = (np.abs(psi_.H * target_psi) ** 2).item(0).real 
        
        psi_list.append(psi_)
        fid_list.append(fid)
    
    if F < max(fid_list):
        best_action = fid_list.index(max(fid_list))
        best_fid = max(fid_list)
    else:
        
        best_action = fid_list.index(min(fid_list))
        best_fid = min(fid_list)
    psi_ = psi_list[best_action]
    # print(best_action)
    return best_action, best_fid, psi_

#动作选最优的，或者次优的
def step2(psi,target_psi,F):
    fid_list = []
    psi_list = []
    action_list = list(range(len(action_space)))
    for action in action_list:
        
        J_1, J_2=  action_space[action,0], action_space[action,1]  # control field strength
        J_12 = J_1 * J_2 /2
    
        H =  (J_1*np.kron(sz, I) + J_2*np.kron(I, sz) + \
                        J_12/2*np.kron((sz - I),(sz - I)) + \
           h_1*np.kron(sx,I) + h_2*np.kron(I,sx))*coupling
        U = expm(-1j * H * dt) 
        psi_ = U * psi
        fid = (np.abs(psi_.H * target_psi) ** 2).item(0).real 
        
        psi_list.append(psi_)
        fid_list.append(fid)
        
    if F < max(fid_list):
        best_action = fid_list.index(max(fid_list))
        best_fid = max(fid_list)
    else:
        psi_list_ = copy.deepcopy(psi_list)
        fid_list_ = copy.deepcopy(fid_list)
        
        del psi_list_[fid_list_.index(max(fid_list_))]
        del fid_list_[fid_list_.index(max(fid_list_))]
        
        best_action = fid_list.index(max(fid_list_))
        
        best_fid = max(fid_list_)
        
        # best_action = fid_list.index(min(fid_list))
        # best_fid = min(fid_list)
    psi_ = psi_list[best_action]
    # print(best_action)
    return best_action, best_fid, psi_

#---------------------------------------------------------------------------------
#将数值列表从小到大进行排序（冒泡法）
def sort_fid_list(lt):
    n= len(lt)
    for x in range(n-1):
       for y in range(n-1-x):
          if lt[y]>lt[y+1]:
             lt[y],lt[y+1]=lt[y+1],lt[y]
    return lt
#--------------------------------------------------------------------------------


time1 = time()
def job(target_psi):
    
    fid_list_noise = [] # 中间元素为列表，每个列表元素为[a,b,c,d,...],其中 a 为无噪声时， 其他为有噪声时
    target_psi = np.mat(target_psi).T
    
    for k in range(init_set_num):
        psi1 = np.mat(init_set[k]).T
        
        fid_list = [] # 元素为[a,b,c,d,...],其中 a 为无噪声时， 其他为有噪声时
        
        action_list_1 = []
        action_list_2 = []
        action_list_0 = []
        
        psi = psi1
        F = (np.abs(psi1.H * target_psi) ** 2).item(0).real 
        
        fid_max1 = F
        fid_max2 = F
        fid_max0 = F
        
        fid_list_1 = [F]
        fid_list_2 = [F]
        fid_list_0 = [F]
        
        step_n = 0
        while True:
            action, F, psi_ = step1(psi,target_psi,F)
            action_list_1.append(action)
            fid_list_1.append(F)
            fid_max1 = max(F,fid_max1)
            psi = psi_
            step_n += 1
            if fid_max1>0.999 or step_n>step_max:
                break
            
        step_n = 0
        F = (np.abs(psi1.H * target_psi) ** 2).item(0).real 
        psi = psi1
        while True:
            action, F, psi_ = step2(psi,target_psi,F)
            action_list_2.append(action)
            fid_list_2.append(F)
            fid_max2 = max(F,fid_max2)
            psi = psi_
            step_n += 1
            if fid_max2>0.999 or step_n>step_max:
                break 
            
        step_n = 0
        F = (np.abs(psi1.H * target_psi) ** 2).item(0).real 
        psi = psi1
        while True:
            action, F, psi_ = step0(psi,target_psi,F)
            action_list_0.append(action)
            fid_list_0.append(F)
            fid_max0 = max(F,fid_max0)
            psi = psi_
            step_n += 1
            if fid_max0>0.999 or step_n>step_max:
                break 
            
        fid_max = max(fid_max0,fid_max1,fid_max2) 
        fid_list.append(fid_max)
        fid_max_index = [fid_max1,fid_max2,fid_max0].index(max([fid_max1,fid_max2,fid_max0]))
        
        if fid_max_index == 0: 
            actions_list = action_list_1
            fids_list = fid_list_1
        elif fid_max_index == 1:
            actions_list = action_list_2
            fids_list = fid_list_2
        elif fid_max_index == 2:
            actions_list = action_list_0
            fids_list = fid_list_0
        
        max_index = fids_list.index(fid_max)
        actions_list = actions_list[0:max_index+1]
        
        #在对应于最大保真度的动作策略上施加噪声
         
        for noise_am in noise_am_list:
            
            noise_list = noise_am * noise_drift 
            # noise_list = noise_am * noise_normal
            psi = psi1
            F = (np.abs(psi1.H * target_psi) ** 2).item(0).real 
            fid_noise_list = [F]
            i = 0
            for action in actions_list:
                noise_unit = [noise_list[i],noise_list[i+1]]
                # F, psi_ =  step_noise_J(psi, target_psi, action, noise_unit)
                F, psi_ =  step_noise_h(psi, target_psi, action, noise_unit)
                i += 2
                fid_noise_list.append(F)
                psi = psi_
            fid_noise = fid_noise_list[max_index]
            fid_list.append(fid_noise)
            
        fid_list_noise.append(fid_list)
        
    fid_list_noise = np.array(fid_list_noise)
            
    return np.mean(fid_list_noise,axis=0)

# noise_am_list = [0.00,0.01,0.02,0.03,0.04,0.05,0.06,0.07,0.08,0.09,0.10]
noise_am_list = [-0.020,-0.018,-0.016,-0.014,-0.012,-0.010,-0.008,-0.006,-0.004,-0.002,0.000,0.002,0.004,0.006,0.008,0.010,0.012,0.014,0.016,0.018,0.020]


def multicore():
    pool = mp.Pool()
    F_list_noise = pool.map(job, target_set)
    return np.mean(np.array(F_list_noise),axis=0)
    
if __name__ == '__main__':

    
    time1 = time()
    
    F_list_noise =  multicore()
    
    print('\n F_list_noise = \n',F_list_noise) 
    
    time2 = time()
    
    print('\n time_cost is: ',time2-time1)
    
# normal

# step_noise_J
# noise_am_list = [0.00,0.01,0.02,0.03,0.04,0.05,0.06,0.07,0.08,0.09,0.10]
 #    F_list_noise = 
  # [0.91085202 0.91085202 0.90879453 0.9025663  0.89229842 0.8782084 0.86059292 0.83981769 0.81630507 0.79052007 0.76295511 0.73411453]

 # time_cost is:  23335.533771276474
 
 # step_noise_h
  #  F_list_noise = 
 # [0.91085202 0.91085202 0.91053506 0.90957449 0.90798189 0.90577077 0.90295651 0.89955614 0.89558824 0.89107272 0.88603073 0.88048441]

 # time_cost is:  21432.222687721252

 
 # drift 
 
 # step_noise_J
 # noise_am_list = [-0.020,-0.018,-0.016,-0.014,-0.012,-0.010,-0.008,-0.006,-0.004,-0.002,0.000,0.002,0.004,0.006,0.008,0.010,0.012,0.014,0.016,0.018,0.020]

 # F_list_noise = 
  # [0.91085202 0.88747087 0.89179427 0.89571301 0.89921057 0.90227189 0.90488349 0.90703364 0.90871246 0.90991204 0.91062648 0.91085202 0.91058705 0.9098321  0.90858994 0.90686544 0.90466563 0.90199955 0.89887823 0.89531456 0.89132314 0.88692019]

 # time_cost is:  23167.32023882866
 
 #  # step_noise_h
 #  F_list_noise = 
 # [0.91085202 0.90494162 0.90612911 0.90718696 0.90811389 0.90890871 0.90957037 0.9100979  0.91049049 0.91074741 0.91086808 0.91085202 0.9106989  0.9104085  0.90998073 0.90941561 0.90871333 0.90787417 0.90689856 0.90578704 0.90454031 0.90315916]

 # time_cost is:  22973.682391881943


